# -*- coding: utf-8 -*-
"""
JSON & BSON codecs and the SocketQueue class which uses them.
"""
import asyncio
from socket import error as socket_error
from struct import unpack

from aiobsonrpc.exceptions import (JsonRpcError, DecodingError, EncodingError, FramingError)

__license__ = 'http://mozilla.org/MPL/2.0/'


class JSONCodec(object):
    """
    Encode/Decode messages to/from JSON format.
    """

    def __init__(self, extractor, framer, custom_codec_implementation=None):
        self._extractor = extractor  # 解析数据流
        self._framer = framer  # 构建数据流 b''
        if custom_codec_implementation is not None:
            self._loads = custom_codec_implementation.loads
            self._dumps = custom_codec_implementation.dumps
        else:
            import json
            self._loads = json.loads
            self._dumps = json.dumps

    def loads(self, b_msg):
        try:
            return self._loads(b_msg.decode('utf-8'))
        except Exception as e:
            raise DecodingError(e)

    def dumps(self, msg):
        try:
            return self._dumps(msg,
                               separators=(',', ':'),
                               sort_keys=True).encode('utf-8')
        except Exception as e:
            raise EncodingError(e)

    def extract_message(self, raw_bytes):
        """
        获取数据流
        """
        try:
            return self._extractor(raw_bytes)
        except Exception as e:
            raise FramingError(e)

    def into_frame(self, message_bytes):
        """
        构建数据流 b''
        """
        try:
            return self._framer(message_bytes)
        except Exception as e:
            raise FramingError(e)


class SocketQueue(object):
    """
    SocketQueue is a duplex Queue connected to a given socket and
    internally takes care of the conversion chain:

    python-data <-> queue-interface <-> codec <-> socket <-:net:-> peer node.
    """

    BUFSIZE = 1024

    SHUT_RDWR = 2

    def __init__(self, reader, writer, loop, codec):
        """
        :param codec: Codec converting python data to/from binary data
        :type codec: BSONCodec or JSONCodec
        """
        if not loop:
            loop = asyncio.get_event_loop()
        self._reader = reader
        self._writer = writer
        self.codec = codec
        self._queue = asyncio.Queue()
        self._lock = asyncio.Lock()
        self._receiver_task = loop.create_task(self._receiver())
        self._closed = False

    @property
    def is_closed(self):
        """
        :property: bool -- Closed by peer node or with ``close()``
        """
        return self._closed

    async def close(self):
        """
        Close this queue and the underlying socket.
        """
        if not self._closed:
            self._writer.close()
            await self._writer.wait_closed()
            self._closed = True

    async def put(self, item):
        """
        Put item to queue -> codec -> :net:

        :param item: Message object.
        :type item: dict, list or None
        """
        if self._closed:
            raise JsonRpcError('Attempt to put items to closed queue.')
        msg_bytes = self.codec.into_frame(self.codec.dumps(item))
        async with self._lock:
            self._writer.write(msg_bytes)
            await self._writer.drain()

    async def get(self):
        """
        Get message items  <- codec <- socket.

        :returns: Normally a message object (python dict or list) but
                  if socket is closed by peer and queue is drained then
                  ``None`` is returned.
                  May also be Exception object in case of parsing or
                  framing errors.
        """
        return await self._queue.get()

    async def _to_queue(self, buffer):
        """
        将接收到的数据写入queue
        """
        b_msg, buffer = self.codec.extract_message(buffer)
        while b_msg is not None:
            await self._queue.put(self.codec.loads(b_msg))
            b_msg, buffer = self.codec.extract_message(buffer)
        return buffer

    async def _receiver(self):
        """
        接收对端数据流处理 -> _to_queue
        """
        buffer = b''
        while True:
            try:
                chunk = await self._reader.read(self.BUFSIZE)
                buffer = await self._to_queue(buffer + chunk)
                if chunk == b'':
                    break
            except DecodingError as e:
                await self._queue.put(e)
            except (OSError, socket_error) as e:
                if e.errno != 9:
                    await self._queue.put(e)
                break
            except Exception as e:
                await self._queue.put(e)
                break
        self._closed = True
        await self._queue.put(None)
        try:  # Just in case somehow socket is still open:
            self._writer.close()
            await self._writer.wait_closed()
        except:
            pass  # Probably already was shut down.

    async def join(self):
        """
        Wait for internal socket receiver thread to finish.
        """
        await self._receiver_task
